package com.ideaaedi.zoo.diy.artifact.websocket.core;

import com.ideaaedi.zoo.commonbase.constant.ZooConstant;
import com.ideaaedi.zoo.commonbase.injector.ExtAutowiredInjector;
import com.ideaaedi.zoo.commonbase.support.AppInstanceIdSupport;
import com.ideaaedi.zoo.diy.artifact.websocket.distribute.DistributedWebsocketMsgConsumer;
import com.ideaaedi.zoo.diy.artifact.websocket.distribute.DistributedWebsocketMsgDTO;
import com.ideaaedi.zoo.diy.artifact.websocket.distribute.DistributedWebsocketMsgProvider;
import jakarta.websocket.OnClose;
import jakarta.websocket.OnError;
import jakarta.websocket.OnMessage;
import jakarta.websocket.OnOpen;
import jakarta.websocket.Session;
import lombok.extern.slf4j.Slf4j;
import org.apache.commons.lang3.BooleanUtils;
import org.apache.commons.lang3.StringUtils;
import org.apache.commons.lang3.tuple.Pair;
import org.slf4j.MDC;
import org.springframework.beans.factory.annotation.Autowired;
import org.springframework.lang.NonNull;

import javax.annotation.Nonnull;
import java.io.IOException;
import java.nio.ByteBuffer;
import java.util.HashSet;
import java.util.List;
import java.util.Map;
import java.util.Optional;
import java.util.Set;
import java.util.concurrent.ConcurrentHashMap;

/**
 * (non-javadoc)
 *
 * @author <font size = "20" color = "#3CAA3C"><a href="https://gitee.com/JustryDeng">JustryDeng</a></font> <img
 * src="https://gitee.com/JustryDeng/shared-files/raw/master/JustryDeng/avatar.jpg" />
 * @since 1.0.0
 */
@Slf4j
public abstract class AbstractWebSocketEndpoint extends ExtAutowiredInjector implements DistributedWebsocketMsgConsumer {
    
    /**
     * 已经移除过本地session的标识
     */
    public static final String ALREADY_REMOVED_LOCAL_SESSION = "__WS_ALREADY_REMOVED_LOCAL_SESSION__";
    
    /**
     * ws的用户id
     */
    public static final String WS_USER_ID_FLAG = "__WS_USER_ID__";
    
    /**
     * <pre>
     * key - session唯一key， {@link AbstractWebSocketEndpoint#sessionUuid(Session)}
     * value - session
     * </pre>
     */
    protected static final Map<String, Session> CLIENT_SESSION_MAP = new ConcurrentHashMap<>(64);
    
    /**
     * ping消息
     */
    protected static final ByteBuffer PING_BYTE_BUFFER = ByteBuffer.wrap("ping".getBytes());
    
    @Autowired(required = false)
    protected WebSocketCallback websocketCallback;
    
    @Autowired(required = false)
    protected DistributedWebsocketMsgProvider distributedWebsocketMsgProvider;
    
    /**
     * 获取与session对应的唯一id
     * <pre>
     * 注：同一个session返回的uuid应始终相同
     * </pre>
     */
    @Nonnull
    protected String sessionUuid(@Nonnull Session session) {
        return appInstanceId() + "_" + session.getId();
    }
    
    /**
     * 服务端收到消息
     */
    @OnMessage
    public abstract void onMessage(String message, @Nonnull Session session);
    
    /**
     * 获取当前应用实例id
     */
    public String appInstanceId() {
        return AppInstanceIdSupport.APP_INSTANCE_ID;
    }
    
    /**
     * 连接建立
     */
    @OnOpen
    public void onOpen(@Nonnull Session session) {
        String sessionUuid = sessionUuid(session);
        String path = session.getRequestURI().getPath();
        String userId = Optional.ofNullable(session.getUserProperties())
                .map(mp -> mp.get(WS_USER_ID_FLAG))
                .map(Object::toString)
                .orElse("0");
        String uri = StringUtils.strip(path.replace("/", "."), ".");
        String traceXd = String.join("_", sessionUuid, userId, uri);
        session.getUserProperties().put(ZooConstant.TRACE_XD, traceXd);
        try (MDC.MDCCloseable ignored = MDC.putCloseable(ZooConstant.TRACE_XD, traceXd)) {
            // 获取到url中的参数：ws://www.idea-aedi.com/ws/{k0}?k1=v1&k2=v2
            Map<String, List<String>> requestParameterMap = session.getRequestParameterMap();
            log.info("session created. path -> {}, requestParameterMap -> {}", path, requestParameterMap);
            // 当前连接已存在，关闭
            Session alreadyExistSession = CLIENT_SESSION_MAP.get(sessionUuid);
            if (alreadyExistSession != null) {
                try {
                    alreadyExistSession.close();
                } catch (Exception e) {
                    log.warn("close already-exist-session exception. e.getMessage() -> {}.", e.getMessage());
                }
            }
            CLIENT_SESSION_MAP.put(sessionUuid, session);
            
            // websocketCallback回调
            if (websocketCallback != null) {
                websocketCallback.onPut(sessionUuid, session);
            }
        }
    }
    
    /**
     * 连接关闭
     */
    @OnClose
    public void onClose(Session session) {
        String traceXd = null;
        try {
            traceXd = (String)session.getUserProperties().get(ZooConstant.TRACE_XD);
        } catch (Exception ignore) {
        }
        try (MDC.MDCCloseable ignored = MDC.putCloseable(ZooConstant.TRACE_XD, traceXd)) {
            log.info("session removed.");
            String sessionUuid = sessionUuid(session);
            // 移除本地session
            removeLocalSession(sessionUuid, websocketCallback);
        }
    }
    
    /**
     * 发生错误时
     */
    @OnError
    public void onError(Session session, Throwable e) {
        String traceXd = null;
        try {
            traceXd = (String)session.getUserProperties().get(ZooConstant.TRACE_XD);
        } catch (Exception ignore) {
        }
        try (MDC.MDCCloseable ignored = MDC.putCloseable(ZooConstant.TRACE_XD, traceXd)) {
            boolean alreadyRemoved = false;
            try {
                alreadyRemoved = session.getUserProperties().containsKey(ALREADY_REMOVED_LOCAL_SESSION);
            } catch (Exception ignore) {
            }
            if (alreadyRemoved) {
                return;
            }
            // 如果不是心跳线程调用触发的，那么需要打印异常
            if (BooleanUtils.isNotTrue(WebSocketHeartbeatRunner.HEARTBEAT_RUNNER_CALL.get())) {
                log.error("websocket occur exception.", e);
            }
            // 移除本地session
            String sessionUuid = sessionUuid(session);
            removeLocalSession(sessionUuid, websocketCallback);
        }
    }
    
    @Override
    public void consume(String sender, @Nonnull String sessionUuid, @Nonnull String message) {
        // 如果这条消息就是自己找不到对应session，才广播出去的，那这里就不需要再消费消息了
        if (StringUtils.equals(sender, appInstanceId())) {
            return;
        }
        sendMessageLocal(sessionUuid, message);
    }
    
    /**
     * 向指定session推送消息（支持分布式）
     *
     * @param sessionUuid session唯一key
     * @param message 消息内容
     */
    public void sendMessageDistributed(String sessionUuid, String message) {
        Session session = CLIENT_SESSION_MAP.get(sessionUuid);
        if (session != null) {
            sendMessageLocal(session, message);
            return;
        }
        if (distributedWebsocketMsgProvider != null) {
            distributedWebsocketMsgProvider.provide(
                    DistributedWebsocketMsgDTO.builder()
                            .sender(appInstanceId())
                            .sessionUuid(sessionUuid)
                            .msg(message)
                            .build()
            );
        }
    }
    
    /**
     * 向本地的指定session推送消息
     *
     * @param sessionUuid session唯一key
     * @param message 消息内容
     */
    public void sendMessageLocal(String sessionUuid, @Nonnull String message) {
        Session session = CLIENT_SESSION_MAP.get(sessionUuid);
        if (session == null) {
            return;
        }
        sendMessageLocal(session, message);
    }
    
    /**
     * 向本地指定session推送文本消息
     */
    public void sendMessageLocal(Session session, String message) {
        String sessionUuid = sessionUuid(session);
        String traceXd = MDC.get(ZooConstant.TRACE_XD);
        boolean originNonTraceXd = traceXd == null;
        try {
            if (originNonTraceXd) {
                try {
                    traceXd = (String)session.getUserProperties().get(ZooConstant.TRACE_XD);
                    MDC.put(ZooConstant.TRACE_XD, traceXd);
                } catch (Exception ignore) {
                }
            }
            session.getAsyncRemote().sendText(message);
            ///session.getAsyncRemote().sendText(message, result -> {
            ///    if (result.isOK()) {
            ///        return;
            ///    }
            ///    log.error("send message '{}' to '{}' fail(async callback). ", message, sessionUuid, result.getException());
            ///    /*
            ///     * 如果客户端异常了需要断开会话，那么会触发OnClose、OnError，或者被WebSocketHeartbeatRunner检测到进而释放资源，
            ///     * 一般不需要在这里主动关闭
            ///     */
            ///    removeLocalSession(sessionUuid, websocketCallback);
            ///});
            log.info("send message '{}' to '{}' success.", message, sessionUuid);
        } catch (Exception e) {
            log.error("send message '{}' to '{}' fail. ", message, sessionUuid, e);
            removeLocalSession(sessionUuid, websocketCallback);
        } finally {
            if (originNonTraceXd) {
                MDC.remove(ZooConstant.TRACE_XD);
            }
        }
    }
    
    /**
     * 向本地指定session推送二进制消息
     */
    public void sendMessageLocal(Session session, ByteBuffer messageByteBuffer) {
        String sessionUuid = sessionUuid(session);
        String traceXd = null;
        try {
            traceXd = (String)session.getUserProperties().get(ZooConstant.TRACE_XD);
        } catch (Exception ignore) {
        }
        try (MDC.MDCCloseable ignored = MDC.putCloseable(ZooConstant.TRACE_XD, traceXd)) {
            session.getAsyncRemote().sendBinary(messageByteBuffer);
            ///session.getAsyncRemote().sendBinary(messageByteBuffer, result -> {
            ///    if (result.isOK()) {
            ///        return;
            ///    }
            ///    /*
            ///     * 如果客户端异常了需要断开会话，那么会触发OnClose、OnError，或者被WebSocketHeartbeatRunner检测到进而释放资源，
            ///     * 一般不需要在这里主动关闭
            ///     */
            ///    removeLocalSession(sessionUuid, websocketCallback);
            ///});
            log.info("send binary-message to '{}' success.", sessionUuid);
        } catch (Exception e) {
            log.info("send binary-message to '{}' fail.", sessionUuid, e);
            removeLocalSession(sessionUuid, websocketCallback);
        }
    }
    
    /**
     * 判断连接是否存活
     *
     * @param session ws session
     *
     * @return 异常对象（为null-存活，不为null-未存活）
     */
    public static Throwable ifAlive(Session session) {
        try {
            session.getBasicRemote().sendPing(PING_BYTE_BUFFER);
            return null;
        } catch (Throwable e) {
            return e;
        }
    }
    
    /**
     * 移除本地session
     *
     * @param sessionUuid session唯一key
     * @param websocketCallback ws回调
     */
    public static void removeLocalSession(String sessionUuid, WebSocketCallback websocketCallback) {
        Session removedSession = CLIENT_SESSION_MAP.remove(sessionUuid);
        if (removedSession == null) {
            return;
        }
        Map<String, Object> userProperties = null;
        try {
            userProperties = removedSession.getUserProperties();
        } catch (Exception ignore) {
        }
        if (userProperties == null || userProperties.containsKey(ALREADY_REMOVED_LOCAL_SESSION)) {
            return;
        }
        userProperties.put(ALREADY_REMOVED_LOCAL_SESSION, 1);
        // 先回调
        if (websocketCallback != null) {
            try {
                websocketCallback.onRemove(sessionUuid);
            } catch (Exception e) {
                log.warn("websocketCallback occur exception. sessionUuid -> {}", sessionUuid, e);
            }
        }
        // 再关session
        if (removedSession.isOpen()) {
            try {
                removedSession.close();
            } catch (IOException ignore) {
            }
        }
    }
    
    /**
     * 查看本地websocket连接存活状态
     *
     * @return <ul>
     *     <li>左（@NonNull）- 存活的本地ws连接唯一标识集合</li>
     *     <li>右（@NonNull）- 未存活的本地ws连接唯一标识集合</li>
     * </ul>
     */
    @NonNull
    public static Pair<Set<String>, Set<String>> aliveStatus() {
        Set<String> aliveSessionUuidSet = new HashSet<>();
        Set<String> diedSessionUuidSet = new HashSet<>();
        Set<Map.Entry<String, Session>> entries = CLIENT_SESSION_MAP.entrySet();
        for (Map.Entry<String, Session> entry : entries) {
            Throwable throwable = ifAlive(entry.getValue());
            String sessionUuid = entry.getKey();
            if (throwable == null) {
                aliveSessionUuidSet.add(sessionUuid);
            } else {
                diedSessionUuidSet.add(sessionUuid);
            }
        }
        return Pair.of(aliveSessionUuidSet, diedSessionUuidSet);
    }
    
}
